#include <cuda_runtime.h>
#include <driver_types.h>
#include <vector>
#include "../include/SSS_soc.h"


#include <thrust/device_ptr.h>
#include <thrust/scan.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>


// #include "/usr/local/cuda/include/thrust/system/cuda/detail/cub/cub.cuh"
// #include "/usr/local/cuda/include/thrust/system/cuda/detail/cub/util_allocator.cuh"
// #include "/usr/local/cuda/include/thrust/system/cuda/detail/cub/device/device_scan.cuh"

// #include <cub/cub.cuh>
// #include <cub/util_allocator.cuh>
// #include <cub/device/device_scan.cuh>


//using namespace cub;
//using namespace cub;


//#include <cub/cub.cuh>

#define block_size 64


void print_matrix_S(SSS_IMAT *A)
{
    int A_nnz = A->num_nnzs;
    int A_row = A->num_rows;
    int A_col = A->num_cols;
    
    printf("A_nnz = %d , A_row = %d , A_col = %d \n",A_nnz,A_row,A_col);
    
    FILE *fp_A;
    char pathA[100];

    char local_time_str[128];
    char *timeA = NULL;

    sprintf(pathA,"Strong Matrix/Strength_Matrix_S_m%d_k%d.mtx",A_row,A_col);
    fp_A=fopen(pathA,"a+"); 

    fprintf(fp_A,"%d %d %d\n", A_row, A_col,A_nnz);
    for (int i = 1; i <=A_row; i++)
    {
            for (int j = A->row_ptr[i-1]; j < A->row_ptr[i]; j++)
            {
                //base0(HYPRE)
                //fprintf(fp_A,"%d %d %f\n",i-1,A_j[j],A_data[j]);
                //base1(AMGX)
                fprintf(fp_A,"%d %d %d\n",i,A->col_idx[j]+1,1);
            }
    }
    fclose(fp_A);
}


void __global__ SSS_Create_SoC_kernel(int row,int *row_ptr,int *col_idx,double *val,double *A_dig_val,int *S_col_idx,double strength_threshold,double max_row_sum)
{

    int tid = blockDim.x * blockIdx.x + threadIdx.x ;

    double row_scl, row_sum;
    int begin_row,end_row;

    if (tid <row)
    {
        // Compute row scale and row sum
        //计算行比例和行和
        row_scl = row_sum = 0.0;

        begin_row = row_ptr[tid];
        end_row = row_ptr[tid + 1];

        for (int j = begin_row; j < end_row; j++) {

            // Originally: Not consider positive entries
            // row_sum += Aj[j];
            //规约
            row_sum += fabs(val[j]);

            // Originally: Not consider positive entries
            // row_scl = max(row_scl, -Aj[j]); // smallest negative
            if (col_idx[j] != tid)   //非对角
                row_scl = SSS_max(row_scl, fabs(val[j]));   // largest abs

        }

        // Multiply by the strength threshold
        row_scl *= strength_threshold;

        // Find diagonal entries of S and remove them later
        // 对角元素的col_idx = 0;
        for (int j = begin_row; j < end_row; j++) {
            if (col_idx[j] == tid) {
                S_col_idx[j] = -1;
                break;
            }
        }

        // Mark entire row as weak couplings if strongly diagonal-dominant
        //如果强对角线占优，则将整行标记为弱耦合
        // Originally: Not consider positive entries
        // if ( ABS(row_sum) > max_row_sum * ABS(diag.d[i]) ) {
        if (row_sum < (2 - max_row_sum) * fabs(A_dig_val[tid])) {
            for (int j = begin_row; j < end_row; j++)
                S_col_idx[j] = -1;
        }
        else {
            for (int j = begin_row; j < end_row; j++) {
                if (-val[j] <= row_scl) 
                S_col_idx[j] = -1;      // only n-couplings
            }
        }
    }
}

void __global__ get_S_row_num_inclusive_kernel(int row,int *old_row_ptr,int *old_col_idx,int *new_row_ptr)
{
        int tid = blockDim.x * blockIdx.x + threadIdx.x ;
        int begin_row,end_row;
        int j;

        if(tid ==0)
        {
            new_row_ptr[tid] = 0;
        }

        if(tid<row)
        {
            begin_row = old_row_ptr[tid];
            end_row = old_row_ptr[tid+1];
            
            
            new_row_ptr[tid+1] =0;


            for (j = begin_row; j < end_row; j++)
            {
                if (old_col_idx[j] > -1)
                {
                    //ic_scan
                    new_row_ptr[tid+1]++;    
                    
                }
            }
        }
}

void __global__ get_S_row_num_exclusive_kernel(int row,int *old_row_ptr,int *old_col_idx,int *new_row_ptr)
{
        int tid = blockDim.x * blockIdx.x + threadIdx.x ;
        int begin_row,end_row;
        int j;

        // if(tid ==0)
        // {
        //     new_row_ptr[tid] = 0;
        // }

        if(tid<row)
        {
            begin_row = old_row_ptr[tid];
            end_row = old_row_ptr[tid+1];
            
            //ic_scan
            //new_row_ptr[tid+1] =0;
            new_row_ptr[tid] =0;

            for (j = begin_row; j < end_row; j++)
            {
                if (old_col_idx[j] > -1)
                {
                    //ic_scan
                    // new_row_ptr[tid+1]++;    
                    new_row_ptr[tid]++;    
                }
            }
        }
}

void __global__ get_S_col_idx_kernel(int row,int *old_row_ptr,int *old_col_idx,int *new_row_ptr,int *new_col_idx)
{
    int tid = blockDim.x * blockIdx.x + threadIdx.x ;
    int begin_row,end_row;
    int new_begin,new_end;
    int j;

    if(tid<row)
    {
        begin_row = old_row_ptr[tid];
        end_row = old_row_ptr[tid + 1];

        new_begin = new_row_ptr[tid];
        new_end = new_row_ptr[tid + 1];
        
        int tmp = 0;
        for (j = begin_row; j < end_row; j++)
        {
            if (old_col_idx[j] > -1)
            {
                new_col_idx[new_begin + tmp] = old_col_idx[j];
                tmp++;
            }
        }
    }
}



static void strong_couplings_cuda(SSS_MAT *A, SSS_IMAT *S, SSS_AMG_PARS *pars)
{
    cudaError_t EE;
    const double max_row_sum = pars->max_row_sum;   //maximal row sum parseter          //最大行和分析器

    printf("max_row_sum = %lf \n",max_row_sum);
    
    double epsilon_str = pars->strong_threshold;   //强度阈值    ε_str 
    
    printf("SoC strong_threshold(before)= %lf\n",epsilon_str);
    
    const int row = A->num_rows, col = A->num_cols, row1 = row + 1;
    const int nnz = A->num_nnzs;

    epsilon_str = 0;
    // local variables
    int i, j, begin_row, end_row;
    double row_scl, row_sum;


    // (1)copy the structure of A to S
    S->num_rows = row;
    S->num_cols = col;
    S->num_nnzs = nnz;

   
    EE = cudaMalloc((void **)&S->d_row_ptr, row1* sizeof(int));
    if(EE != cudaSuccess)
    {
        printf("qqqqqqqqqqqqqqq \n");
    }
    EE = cudaMalloc((void **)&S->d_col_idx, nnz* sizeof(int));
    if(EE != cudaSuccess)
    {
        printf("wwwwwwwwwwwwwwww \n");
    }
    EE = cudaMemcpy(S->d_row_ptr,A->d_row_ptr,sizeof(int)*(A->num_rows+1),cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("d_row_ptr fail \n");
    }
    
    EE = cudaMemcpy(S->d_col_idx,A->d_col_idx,sizeof(int)*A->num_nnzs,cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("d_col_idx fail \n");
    }
    

    int blocksize = 64;
    int gridsize = ceil ((double)row / (double)blocksize);


    SSS_Create_SoC_kernel<<<gridsize,blocksize>>>(row,A->d_row_ptr,A->d_col_idx,A->d_val,A->d_dig,S->d_col_idx,epsilon_str,max_row_sum);

    cudaDeviceSynchronize();


    printf("SoC strong_threshold(after)= %lf\n",epsilon_str);
    int *row_num_device = NULL;

    EE = cudaMalloc((void **)&row_num_device, row1* sizeof(int));
    if(EE != cudaSuccess)
    {
        printf("row_num_device fail \n");
    }
    

    get_S_row_num_exclusive_kernel<<<gridsize,blocksize>>>(row,S->d_row_ptr,S->d_col_idx,row_num_device);
    cudaDeviceSynchronize();

  

    // CPU scan
    //int *scan_row = (int *) malloc( (row+1) *   sizeof(int));
    //cudaMemcpy(scan_row,row_num_device,sizeof(int)*(row+1),cudaMemcpyDeviceToHost);
    // thrust::exclusive_scan(scan_row, scan_row +(row+1), scan_row); // in-place scan
    // EE = cudaMemcpy(new_row_ptr_device,scan_row, row1 * sizeof(int), cudaMemcpyHostToDevice);
    // if(EE != cudaSuccess)
    // {
    //     printf("scan_row fail \n");
    // }
    

    //new_row_ptr_device
    int *new_row_ptr_device = NULL;
    EE = cudaMalloc((void **)&new_row_ptr_device, row1* sizeof(int));
    if(EE != cudaSuccess)
    {
        printf("qqqqqqqqqqqqqqq \n");
    }

    // GPU scan
    // void            *d_temp_storage = NULL;
    // size_t          temp_storage_bytes = 0;
    // CubDebugExit(cub::DeviceScan::ExclusiveSum(d_temp_storage, temp_storage_bytes, row_num_device, new_row_ptr_device, row1));
    
    cudaDeviceSynchronize();


     thrust::exclusive_scan(thrust::device,row_num_device, row_num_device +(row1), new_row_ptr_device); // in-place scan
    //cudaMemcpy(new_row_ptr_device,row_num_device,sizeof(int)*(row+1),cudaMemcpyDeviceToDevice);
    //cudaDeviceSynchronize();
 

    //new_col_idx_device
    int *new_col_idx_device = NULL;
    EE = cudaMalloc((void **)&new_col_idx_device, S->num_nnzs* sizeof(int));
    if(EE != cudaSuccess)
    {
        printf("qqqqqqqqqqqqqqq \n");
    }


    get_S_col_idx_kernel<<<gridsize,blocksize>>>(row,S->d_row_ptr,S->d_col_idx,new_row_ptr_device,new_col_idx_device);
    cudaDeviceSynchronize();


    EE = cudaMemcpy(S->d_row_ptr,new_row_ptr_device, (row+1)* sizeof(int), cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("S new row_ptr fail \n");
    }

    EE = cudaMemcpy(&S->num_nnzs,&new_row_ptr_device[row],  sizeof(int), cudaMemcpyDeviceToHost);
    if(EE != cudaSuccess)
    {
        printf("S nnz fail \n");
    }

    EE = cudaMemcpy(S->d_col_idx,new_col_idx_device, S->num_nnzs* sizeof(int), cudaMemcpyDeviceToDevice);
    if(EE != cudaSuccess)
    {
        printf("S new col_idx fail \n");
    }

    //free(scan_row);

    EE = cudaFree(row_num_device);
    if(EE != cudaSuccess)
    {
        printf("row_num_device  Failed \n");
    }

    EE = cudaFree(new_row_ptr_device);   
    if(EE != cudaSuccess)
    {
        printf("new_row_ptr_device  Failed \n");
    }
    
    EE = cudaFree(new_col_idx_device);
    if(EE != cudaSuccess)
    {
        printf("new_col_idx_device  Failed \n");
    }




} 









int SoC_GPU(SSS_MAT *A, SSS_IMAT *S, SSS_AMG_PARS *pars)
{


    //CPU 原版
    //strong_couplings_compress(A,S,pars);


    //GPU
    strong_couplings_cuda(A,S,pars);
    


    return 0;
}
